{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "436accd8",
   "metadata": {},
   "source": [
    "# 自然语言推断：使用注意力\n",
    ":label:`sec_natural-language-inference-attention`\n",
    "\n",
    "我们在 :numref:`sec_natural-language-inference-and-dataset`中介绍了自然语言推断任务和SNLI数据集。鉴于许多模型都是基于复杂而深度的架构，Parikh等人提出用注意力机制解决自然语言推断问题，并称之为“可分解注意力模型” :cite:`Parikh.Tackstrom.Das.ea.2016`。这使得模型没有循环层或卷积层，在SNLI数据集上以更少的参数实现了当时的最佳结果。本节将描述并实现这种基于注意力的自然语言推断方法（使用MLP），如 :numref:`fig_nlp-map-nli-attention`中所述。\n",
    "\n",
    "![将预训练GloVe送入基于注意力和MLP的自然语言推断架构](../img/nlp-map-nli-attention.svg)\n",
    ":label:`fig_nlp-map-nli-attention`\n",
    "\n",
    "## 模型\n",
    "\n",
    "与保留前提和假设中词元的顺序相比，我们可以将一个文本序列中的词元与另一个文本序列中的每个词元对齐，然后比较和聚合这些信息，以预测前提和假设之间的逻辑关系。与机器翻译中源句和目标句之间的词元对齐类似，前提和假设之间的词元对齐可以通过注意力机制灵活地完成。\n",
    "\n",
    "![利用注意力机制进行自然语言推断](../img/nli-attention.svg)\n",
    ":label:`fig_nli_attention`\n",
    "\n",
    " :numref:`fig_nli_attention`描述了使用注意力机制的自然语言推断方法。从高层次上讲，它由三个联合训练的步骤组成：对齐、比较和汇总。我们将在下面一步一步地对它们进行说明。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "908a15ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00b36e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "from d2l import mindspore as d2l\n",
    "from mindspore import nn, ops, Tensor\n",
    "from mindspore import dtype as mstype"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c59118e6",
   "metadata": {},
   "source": [
    "### 注意（Attending）\n",
    "\n",
    "第一步是将一个文本序列中的词元与另一个序列中的每个词元对齐。假设前提是“我确实需要睡眠”，假设是“我累了”。由于语义上的相似性，我们不妨将假设中的“我”与前提中的“我”对齐，将假设中的“累”与前提中的“睡眠”对齐。同样，我们可能希望将前提中的“我”与假设中的“我”对齐，将前提中的“需要”和“睡眠”与假设中的“累”对齐。请注意，这种对齐是使用加权平均的“软”对齐，其中理想情况下较大的权重与要对齐的词元相关联。为了便于演示， :numref:`fig_nli_attention`以“硬”对齐的方式显示了这种对齐方式。\n",
    "\n",
    "现在，我们更详细地描述使用注意力机制的软对齐。用$\\mathbf{A} = (\\mathbf{a}_1, \\ldots, \\mathbf{a}_m)$和$\\mathbf{B} = (\\mathbf{b}_1, \\ldots, \\mathbf{b}_n)$表示前提和假设，其词元数量分别为$m$和$n$，其中$\\mathbf{a}_i, \\mathbf{b}_j \\in \\mathbb{R}^{d}$（$i = 1, \\ldots, m, j = 1, \\ldots, n$）是$d$维的词向量。对于软对齐，我们将注意力权重$e_{ij} \\in \\mathbb{R}$计算为：\n",
    "\n",
    "$$e_{ij} = f(\\mathbf{a}_i)^\\top f(\\mathbf{b}_j),$$\n",
    ":eqlabel:`eq_nli_e`\n",
    "\n",
    "其中函数$f$是在下面的`mlp`函数中定义的多层感知机。输出维度$f$由`mlp`的`num_hiddens`参数指定。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12571d03",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(num_inputs, num_hiddens, flatten):\n",
    "    net = []\n",
    "    net.append(nn.Dropout(p=0.2))\n",
    "    net.append(nn.Dense(num_inputs, num_hiddens))\n",
    "    net.append(nn.ReLU())\n",
    "    if flatten:\n",
    "        net.append(nn.Flatten()) \n",
    "    net.append(nn.Dropout(p=0.2))\n",
    "    net.append(nn.Dense(num_hiddens, num_hiddens))\n",
    "    net.append(nn.ReLU())\n",
    "    if flatten:\n",
    "        net.append(nn.Flatten())\n",
    "    return nn.SequentialCell(*net)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36a59699",
   "metadata": {},
   "source": [
    "值得注意的是，在 :eqref:`eq_nli_e`中，$f$分别输入$\\mathbf{a}_i$和$\\mathbf{b}_j$，而不是将它们一对放在一起作为输入。这种*分解*技巧导致$f$只有$m + n$个次计算（线性复杂度），而不是$mn$次计算（二次复杂度）\n",
    "\n",
    "对 :eqref:`eq_nli_e`中的注意力权重进行规范化，我们计算假设中所有词元向量的加权平均值，以获得假设的表示，该假设与前提中索引$i$的词元进行软对齐：\n",
    "\n",
    "$$\n",
    "\\boldsymbol{\\beta}_i = \\sum_{j=1}^{n}\\frac{\\exp(e_{ij})}{ \\sum_{k=1}^{n} \\exp(e_{ik})} \\mathbf{b}_j.\n",
    "$$\n",
    "\n",
    "同样，我们计算假设中索引为$j$的每个词元与前提词元的软对齐：\n",
    "\n",
    "$$\n",
    "\\boldsymbol{\\alpha}_j = \\sum_{i=1}^{m}\\frac{\\exp(e_{ij})}{ \\sum_{k=1}^{m} \\exp(e_{kj})} \\mathbf{a}_i.\n",
    "$$\n",
    "\n",
    "下面，我们定义`Attend`类来计算假设（`beta`）与输入前提`A`的软对齐以及前提（`alpha`）与输入假设`B`的软对齐。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d2e811f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attend(nn.Cell):\n",
    "    def __init__(self, num_inputs, num_hiddens, **kwargs):\n",
    "        super(Attend, self).__init__(**kwargs)\n",
    "        self.f = mlp(num_inputs, num_hiddens, flatten=False)\n",
    "\n",
    "    def construct(self, A, B):\n",
    "        # A/B的形状：（批量大小，序列A/B的词元数，embed_size）\n",
    "        # f_A/f_B的形状：（批量大小，序列A/B的词元数，num_hiddens）\n",
    "        f_A = self.f(A)\n",
    "        f_B = self.f(B)\n",
    "        # e的形状：（批量大小，序列A的词元数，序列B的词元数）\n",
    "        e = f_A.bmm(f_B.transpose(0, 2, 1))\n",
    "        # beta的形状：（批量大小，序列A的词元数，embed_size），\n",
    "        # 意味着序列B被软对齐到序列A的每个词元(beta的第1个维度)\n",
    "        beta = ops.softmax(e, axis=-1).bmm(B)\n",
    "        # beta的形状：（批量大小，序列B的词元数，embed_size），\n",
    "        # 意味着序列A被软对齐到序列B的每个词元(alpha的第1个维度)\n",
    "        alpha = ops.softmax(e.transpose(0, 2, 1), axis=-1).bmm(A)\n",
    "        return beta, alpha"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff3c583",
   "metadata": {},
   "source": [
    "### 比较\n",
    "\n",
    "在下一步中，我们将一个序列中的词元与与该词元软对齐的另一个序列进行比较。请注意，在软对齐中，一个序列中的所有词元（尽管可能具有不同的注意力权重）将与另一个序列中的词元进行比较。为便于演示， :numref:`fig_nli_attention`对词元以*硬*的方式对齐。例如，上述的*注意*（attending）步骤确定前提中的“need”和“sleep”都与假设中的“tired”对齐，则将对“疲倦-需要睡眠”进行比较。\n",
    "\n",
    "在比较步骤中，我们将来自一个序列的词元的连结（运算符$[\\cdot, \\cdot]$）和来自另一序列的对齐的词元送入函数$g$（一个多层感知机）：\n",
    "\n",
    "$$\\mathbf{v}_{A,i} = g([\\mathbf{a}_i, \\boldsymbol{\\beta}_i]), i = 1, \\ldots, m\\\\ \\mathbf{v}_{B,j} = g([\\mathbf{b}_j, \\boldsymbol{\\alpha}_j]), j = 1, \\ldots, n.$$\n",
    "\n",
    ":eqlabel:`eq_nli_v_ab`\n",
    "\n",
    "在 :eqref:`eq_nli_v_ab`中，$\\mathbf{v}_{A,i}$是指，所有假设中的词元与前提中词元$i$软对齐，再与词元$i$的比较；而$\\mathbf{v}_{B,j}$是指，所有前提中的词元与假设中词元$i$软对齐，再与词元$i$的比较。下面的`Compare`个类定义了比较步骤。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fabac37",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Compare(nn.Cell):\n",
    "    def __init__(self, num_inputs, num_hiddens, **kwargs):\n",
    "        super(Compare, self).__init__(**kwargs)\n",
    "        self.g = mlp(num_inputs, num_hiddens, flatten=False)\n",
    "\n",
    "    def construct(self, A, B, beta, alpha):\n",
    "        V_A = self.g(ops.cat((A, beta), axis=2))\n",
    "        V_B = self.g(ops.cat((B, alpha), axis=2))\n",
    "        return V_A, V_B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6022fa11",
   "metadata": {},
   "source": [
    "### 聚合\n",
    "\n",
    "现在我们有两组比较向量$\\mathbf{v}_{A,i}$（$i = 1, \\ldots, m$）和$\\mathbf{v}_{B,j}$（$j = 1, \\ldots, n$）。在最后一步中，我们将聚合这些信息以推断逻辑关系。我们首先求和这两组比较向量：\n",
    "\n",
    "$$\n",
    "\\mathbf{v}_A = \\sum_{i=1}^{m} \\mathbf{v}_{A,i}, \\quad \\mathbf{v}_B = \\sum_{j=1}^{n}\\mathbf{v}_{B,j}.\n",
    "$$\n",
    "\n",
    "接下来，我们将两个求和结果的连结提供给函数$h$（一个多层感知机），以获得逻辑关系的分类结果：\n",
    "\n",
    "$$\n",
    "\\hat{\\mathbf{y}} = h([\\mathbf{v}_A, \\mathbf{v}_B]).\n",
    "$$\n",
    "\n",
    "聚合步骤在以下`Aggregate`类中定义。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b92c515",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Aggregate(nn.Cell):\n",
    "    def __init__(self, num_inputs, num_hiddens, num_outputs, **kwargs):\n",
    "        super(Aggregate, self).__init__(**kwargs)\n",
    "        self.h = mlp(num_inputs, num_hiddens, flatten=True)\n",
    "        self.linear = nn.Dense(num_hiddens, num_outputs)\n",
    "\n",
    "    def construct(self, V_A, V_B):\n",
    "        # 对两组比较向量分别求和\n",
    "        V_A = V_A.sum(axis=1)\n",
    "        V_B = V_B.sum(axis=1)\n",
    "        # 将两个求和结果的连结送到多层感知机中\n",
    "        Y_hat = self.linear(self.h(ops.cat((V_A, V_B), axis=1)))\n",
    "        return Y_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c1835ba",
   "metadata": {},
   "source": [
    "### 整合代码\n",
    "\n",
    "通过将注意步骤、比较步骤和聚合步骤组合在一起，我们定义了可分解注意力模型来联合训练这三个步骤。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f725dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecomposableAttention(nn.Cell):\n",
    "    def __init__(self, vocab, embed_size, num_hiddens, num_inputs_attend=100,\n",
    "                 num_inputs_compare=200, num_inputs_agg=400, **kwargs):\n",
    "        super(DecomposableAttention, self).__init__(**kwargs)\n",
    "        self.embedding = nn.Embedding(len(vocab), embed_size)\n",
    "        self.attend = Attend(num_inputs_attend, num_hiddens)\n",
    "        self.compare = Compare(num_inputs_compare, num_hiddens)\n",
    "        # 有3种可能的输出：蕴涵、矛盾和中性\n",
    "        self.aggregate = Aggregate(num_inputs_agg, num_hiddens, num_outputs=3)\n",
    "\n",
    "    def construct(self, X):\n",
    "        premises, hypotheses = X[:, 0], X[:, 1]\n",
    "        A = self.embedding(premises)\n",
    "        B = self.embedding(hypotheses)\n",
    "        beta, alpha = self.attend(A, B)\n",
    "        V_A, V_B = self.compare(A, B, beta, alpha)\n",
    "        Y_hat = self.aggregate(V_A, V_B)\n",
    "        return Y_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7633ebc8",
   "metadata": {},
   "source": [
    "## 训练和评估模型\n",
    "\n",
    "现在，我们将在SNLI数据集上对定义好的可分解注意力模型进行训练和评估。我们从读取数据集开始。\n",
    "\n",
    "### 读取数据集\n",
    "\n",
    "我们使用 :numref:`sec_natural-language-inference-and-dataset`中定义的函数下载并读取SNLI数据集。批量大小和序列长度分别设置为$256$和$50$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e1fbafb",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size, num_steps = 256, 50\n",
    "train_iter, test_iter, vocab = d2l.load_data_snli(batch_size, num_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14391a85",
   "metadata": {},
   "source": [
    "### 创建模型\n",
    "\n",
    "我们使用预训练好的100维GloVe嵌入来表示输入词元。我们将向量$\\mathbf{a}_i$和$\\mathbf{b}_j$在 :eqref:`eq_nli_e`中的维数预定义为100。 :eqref:`eq_nli_e`中的函数$f$和 :eqref:`eq_nli_v_ab`中的函数$g$的输出维度被设置为200.然后我们创建一个模型实例，初始化它的参数，并加载GloVe嵌入来初始化输入词元的向量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89a270df",
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_size, num_hiddens = 100, 200\n",
    "net = DecomposableAttention(vocab, embed_size, num_hiddens)\n",
    "glove_embedding = d2l.TokenEmbedding('glove.6b.100d')\n",
    "embeds = glove_embedding[vocab.idx_to_token]\n",
    "net.embedding.embedding_table.set_data(d2l.tensor(embeds))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83e22926",
   "metadata": {},
   "source": [
    "### 训练和评估模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e86895a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr, num_epochs = 0.001, 4\n",
    "trainer = nn.Adam(learning_rate=lr, params=net.trainable_params())\n",
    "loss = nn.CrossEntropyLoss(reduction=\"none\")\n",
    "d2l.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9a6d5e2",
   "metadata": {},
   "source": [
    "### 使用模型\n",
    "\n",
    "最后，定义预测函数，输出一对前提和假设之间的逻辑关系。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2f40a3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_snli(net, vocab, premise, hypothesis):\n",
    "    \"\"\"预测前提和假设之间的逻辑关系\"\"\"\n",
    "    net.set_train(False)\n",
    "    premise = Tensor(vocab[premise], dtype=mstype.int32)\n",
    "    hypothesis = Tensor(vocab[hypothesis], dtype=mstype.int32)\n",
    "    X = ops.stack((premise.reshape(-1, 1), hypothesis.reshape(-1, 1)), axis=1)\n",
    "    label = ops.argmax(net(X), axis=-1)\n",
    "    results = []\n",
    "    for l in label:\n",
    "        if l == 0:\n",
    "            results.append('entailment')\n",
    "        elif l == 1:\n",
    "            results.append('contradiction')\n",
    "        else:\n",
    "            results.append(\"neutral\")\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edb82c07",
   "metadata": {},
   "source": [
    "我们可以使用训练好的模型来获得对示例句子的自然语言推断结果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4f3d6c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_snli(net, vocab, ['he', 'is', 'good', '.'], ['he', 'is', 'bad', '.'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50b9e0d9",
   "metadata": {},
   "source": [
    "## 小结\n",
    "\n",
    "* 可分解注意模型包括三个步骤来预测前提和假设之间的逻辑关系：注意、比较和聚合。\n",
    "* 通过注意力机制，我们可以将一个文本序列中的词元与另一个文本序列中的每个词元对齐，反之亦然。这种对齐是使用加权平均的软对齐，其中理想情况下较大的权重与要对齐的词元相关联。\n",
    "* 在计算注意力权重时，分解技巧会带来比二次复杂度更理想的线性复杂度。\n",
    "* 我们可以使用预训练好的词向量作为下游自然语言处理任务（如自然语言推断）的输入表示。\n",
    "\n",
    "## 练习\n",
    "\n",
    "1. 使用其他超参数组合训练模型，能在测试集上获得更高的准确度吗？\n",
    "1. 自然语言推断的可分解注意模型的主要缺点是什么？\n",
    "1. 假设我们想要获得任何一对句子的语义相似级别（例如，0～1之间的连续值）。我们应该如何收集和标注数据集？请尝试设计一个有注意力机制的模型。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
